from tensorflow.python.platform import gfile
import tensorflow as tf
import os
import cv2
import numpy as np
from utiles.box_utiles import polygons_from_bitmap
from utiles.transform import resize_image

label2num_dic = {"0": 0,
              "1": 1,
              "2": 2,
              "3": 3,
              "4": 4,
              "5": 5,
              "6": 6,
              "7": 7,
              "8": 8,
              "9": 9,
              "A": 10,
              "B": 11,
              "C": 12,
              "D": 13,
              "E": 14,
              "F": 15,
              "G": 16,
              "H": 17,
              "J": 18,
              "K": 19,
              "P": 20,
              "S": 21,
              "R": 22,
              "Q": 23}

num2label_dic = {"0": "0",
              "1": "1",
              "2": "2",
              "3": "3",
              "4": "4",
              "5": "5",
              "6": "6",
              "7": "7",
              "8": "8",
              "9": "9",
              "10": "A",
              "11": "B",
              "12": "C",
              "13": "D",
              "14": "E",
              "15": "F",
              "16": "G",
              "17": "H",
              "18": "J",
              "19": "K",
              "20": "P",
              "21": "S",
              "22": "R",
              "23": "Q"}

IMAGE_SIZE = [160, 160]
CUT_SIZE = (32, 32)
db_model_path = './pbMode/db_model.pb'
dec_model_path = './pbMode/dec_model.pb'
image_root = 'E:/DATA/ocr_img/cut_image/'

sess = tf.Session()
with gfile.FastGFile(db_model_path, 'rb') as f:
    graph_def = tf.GraphDef()
    graph_def.ParseFromString(f.read())
    sess.graph.as_default()
    tf.import_graph_def(graph_def, name='')

with gfile.FastGFile(dec_model_path, 'rb') as f:
    graph_def = tf.GraphDef()
    graph_def.ParseFromString(f.read())
    sess.graph.as_default()
    tf.import_graph_def(graph_def, name='')

sess.run(tf.global_variables_initializer())

db_input_image = sess.graph.get_tensor_by_name('image_input:0')
db_out = sess.graph.get_tensor_by_name('dbnet/proba3_sigmoid:0')
dec_input_image = sess.graph.get_tensor_by_name('Image:0')
decision_out = sess.graph.get_tensor_by_name('decision_out:0')
image_names = [i[2] for i in os.walk(image_root)][0]
name_dic = {}

for i in image_names:
    image_path = os.path.join(image_root, i)
    image = cv2.imread(image_path)
    image = image.astype(np.float32)

    image = resize_image(IMAGE_SIZE[0], image)
    image = np.array(image[np.newaxis,:, :, :])
    # mean = [103.939, 116.779, 123.68]
    # image[..., 0] -= mean[0]
    # image[..., 1] -= mean[1]
    # image[..., 2] -= mean[2]
    image = image/255.0
    proba_map = sess.run(db_out, feed_dict={db_input_image: image})

    proba_map = np.squeeze(proba_map, 0)
    image = np.squeeze(image, 0)
    #
    # image[..., 0] += mean[0]
    # image[..., 1] += mean[1]
    # image[..., 2] += mean[2]
    image *= 255.0
    bitmap = proba_map > 0.3
    boxes, scores = polygons_from_bitmap(proba_map, bitmap, image.shape[0], image.shape[1], box_thresh=0.5)

    name_count = 0
    for contour in boxes:
        contour = np.reshape(contour, (-1, 2))
        xmin = sorted(list(contour), key=lambda x: x[0])[0][0]
        xmax = sorted(list(contour), key=lambda x: x[0])[-1][0]
        ymin = sorted(list(contour), key=lambda x: x[1])[0][1]
        ymax = sorted(list(contour), key=lambda x: x[1])[-1][1]
        cut_image = image[ymin:ymax, xmin:xmax, 0]
        cut_image = cv2.resize(cut_image, CUT_SIZE)
        # cv2.imshow('cut_image', cut_image/255.0)
        # cv2.waitKey()

        cut_image = np.array(cut_image[np.newaxis, :, :, np.newaxis])
        decision = sess.run(decision_out, feed_dict={dec_input_image: cut_image/255.0})[0]
        label = num2label_dic[str(decision)]
        if label not in name_dic.keys():
            name_dic[label] = 1
        else:
            name_dic[label] += 1
        cut_image = np.squeeze(cut_image, 0)
        cv2.imwrite(os.path.join('./cut/', label + '-' + str(name_dic[label]) + '.jpg'), cut_image)











